import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import os

def create_heatmap():
    data_dir = './data/'
    all_months_data = []
    months = [f"{i:02d}" for i in range(1, 13)]

    for month in months:
        file_name = f'GenerationbyFuelType_2018{month}.csv'
        file_path = os.path.join(data_dir, file_name)
        if os.path.exists(file_path):
            # 与之前逻辑相同，加载并计算每小时总和
            df = pd.read_csv(file_path, header=None)
            numeric_cols = df.columns[2:]
            df_numeric = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
            half_hourly_sum = df_numeric.sum(axis=1)
            hourly_index = [i // 2 for i in range(len(half_hourly_sum))]
            hourly_sum = half_hourly_sum.groupby(hourly_index).sum()
            hourly_sum.name = month # 设置Series的名字为月份
            all_months_data.append(hourly_sum)

    # 合并所有月份数据到一个DataFrame
    heatmap_data = pd.concat(all_months_data, axis=1).T # .T进行转置，让月份为行
    
    # 绘制热力图
    plt.figure(figsize=(12, 8))
    # 使用中文显示需要设置字体
    plt.rcParams['font.sans-serif'] = ['SimHei'] 
    plt.rcParams['axes.unicode_minus'] = False 
    
    sns.heatmap(heatmap_data, cmap='viridis', annot=False)
    plt.title('2018年每小时平均用电量热力图')
    plt.xlabel('小时')
    plt.ylabel('月份')
    plt.savefig("energy_heatmap_false.png") # 保存结果图
    plt.show()

if __name__ == '__main__':
    create_heatmap()